package com.heima.ai.controller;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.cloud.ai.dashscope.chat.MessageFormat;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.heima.ai.mapper.ChatHistoryMapper;
import com.heima.ai.pojo.ChatHistory;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.ObjectUtils;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Media;
import org.springframework.util.MimeType;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@RestController
@RequestMapping("/ai")
@RequiredArgsConstructor
public class ChatController {

    private final ChatClient ollamaClient;
    private final ChatHistoryMapper chatHistoryMapper;

    //聊天
    @PostMapping(value = "/chat", produces = "text/html;charset=utf-8")
    public Flux<String> chat(String prompt,
                             String chatId,
                             @RequestParam(value = "files", required = false) List<MultipartFile> files
    ) {

        //根据类型和id查询, 如果查不到, 就进行保存
        LambdaQueryWrapper<ChatHistory> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(ChatHistory::getType, "chat").eq(ChatHistory::getChatId, chatId);
        ChatHistory chatHistory = chatHistoryMapper.selectOne(wrapper);
        if (ObjectUtils.isEmpty(chatHistory)) {
            chatHistoryMapper.insert(new ChatHistory(null, "chat", chatId));
        }
        // 2.请求模型
        if (files == null || files.isEmpty()) {
            // 没有附件，纯文本聊天
            return ollamaClient
                    .prompt(prompt) //用户提示词 user
                    .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)) //会话身份象征
                    .stream()//流式调用
                    .content();//取返回内容
        } else {
            // 将上传的文件转换为媒体对象列表（仅处理非空文件）
            // 1.解析多媒体
            List<Media> medias = files.stream()
                    .map(file -> new Media(
                                    MimeType.valueOf(Objects.requireNonNull(file.getContentType())),
                                    file.getResource()
                            )
                    )
                    .toList();
            var userMessage = new UserMessage(prompt, medias);
            var prompt2 = new Prompt(userMessage, DashScopeChatOptions.builder()
                    .withModel("qwen-omni-turbo")//模型
                    .withMultiModel(true)//多模型
                    .build());

            // 2.请求模型
            return ollamaClient.prompt(prompt2)
                    .advisors(a -> a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                    .stream()
                    .content();
        }
    }
}